import pandas as pd
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LogisticRegression, LinearRegression
from sklearn.ensemble import RandomForestClassifier, RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.metrics import accuracy_score, classification_report, mean_squared_error, r2_score
from utils.data_handler import DataHandler

class MLModels:
    def __init__(self, data, target_column, task_type):
        self.data = data
        self.target_column = target_column
        self.task_type = task_type
        self.data_handler = DataHandler(data)
        
    def train_models(self):
        """训练多个模型并返回结果"""
        # 准备数据
        X, y = self.data_handler.prepare_data(self.target_column)
        
        if X.empty:
            raise ValueError("没有可用的数值特征")
        
        # 分割数据
        X_train, X_test, y_train, y_test = train_test_split(
            X, y, test_size=0.2, random_state=42
        )
        
        # 标准化
        scaler = StandardScaler()
        X_train_scaled = scaler.fit_transform(X_train)
        X_test_scaled = scaler.transform(X_test)
        
        results = {
            'feature_names': X.columns.tolist(),
            'scaler': scaler,
            'models': {},
            'test_data': (X_test_scaled, y_test)
        }
        
        if self.task_type == 'classification':
            models = {
                'logistic_regression': LogisticRegression(random_state=42),
                'random_forest': RandomForestClassifier(n_estimators=100, random_state=42)
            }
            
            for name, model in models.items():
                model.fit(X_train_scaled, y_train)
                y_pred = model.predict(X_test_scaled)
                
                results['models'][name] = {
                    'model': model,
                    'accuracy': accuracy_score(y_test, y_pred),
                    'report': classification_report(y_test, y_pred, output_dict=True)
                }
        
        else:  # regression
            models = {
                'linear_regression': LinearRegression(),
                'random_forest': RandomForestRegressor(n_estimators=100, random_state=42)
            }
            
            for name, model in models.items():
                model.fit(X_train_scaled, y_train)
                y_pred = model.predict(X_test_scaled)
                
                results['models'][name] = {
                    'model': model,
                    'mse': mean_squared_error(y_test, y_pred),
                    'r2': r2_score(y_test, y_pred)
                }
        
        # 选择最佳模型
        if self.task_type == 'classification':
            best_model = max(results['models'].items(), key=lambda x: x[1]['accuracy'])
        else:
            best_model = max(results['models'].items(), key=lambda x: x[1]['r2'])
        
        results['models']['best_model'] = best_model[1]
        results['best_model_name'] = best_model[0]
        
        return results